from typing import Dict, Any, List
import networkx as nx
import numpy as np
from bson import ObjectId
from pymongo import MongoClient
from config.config import Config


def build_graph_from_feature(ast_feature: Dict[str, Any]) -> nx.DiGraph:
    """
    将AST特征字典构建成NetworkX DiGraph对象。
    使用批量添加节点和边的方式以提高效率，并增加输入检查。

    Args:
        ast_feature (dict): 包含 'nodes' 和 'edges' 键的AST特征字典。
                           'nodes' 应为包含节点信息的列表。
                           'edges' 应为包含边信息的列表。

    Returns:
        nx.DiGraph: 构建好的有向图。

    Raises:
        ValueError: 如果输入的 ast_feature 格式不正确。
        KeyError: 如果节点或边缺少必要的键。
    """
    # 1. 输入验证
    if not isinstance(ast_feature, dict):
        raise ValueError("ast_feature 必须是一个字典。")
    if "nodes" not in ast_feature or "edges" not in ast_feature:
        raise ValueError("ast_feature 必须包含 'nodes' 和 'edges' 键。")
    if not isinstance(ast_feature["nodes"], list):
        raise ValueError("'nodes' 必须是一个列表。")
    if not isinstance(ast_feature["edges"], list):
        raise ValueError("'edges' 必须是一个列表。")

    # 2. 创建一个空的有向图
    g = nx.DiGraph()

    # --- 优化点1: 批量添加节点和属性 ---
    # 先准备节点ID列表
    node_ids: List[Any] = []
    # 准备节点属性字典
    node_kinds: Dict[Any, Any] = {}
    node_simhashes: Dict[Any, Any] = {}
    node_tokens: Dict[Any, List] = {}
    node_diff_tokens: Dict[Any, List] = {}

    # 遍历节点，收集数据
    # 使用 enumerate 可以在出错时提供更多信息
    try:
        for i, n in enumerate(ast_feature["nodes"]):
            if not isinstance(n, dict) or "id" not in n:
                raise KeyError(f"节点列表中第 {i} 个元素缺少 'id' 键或不是字典: {n}")
            node_id = n["id"]
            node_ids.append(node_id)

            # 收集各个属性，用于批量设置
            # 假设 'type' 键是必须存在的，否则会抛出 KeyError
            node_kinds[node_id] = n["type"]
            # 对于可能不存在且默认值为 None 或 [] 的属性，get 是合适的
            node_tokens[node_id] = n.get("nodeToken", [])  # 默认 []
            node_diff_tokens[node_id] = n.get("diffnochangeToken", [])  # 默认 []
    except KeyError as e:
        raise KeyError(f"处理节点时出错: {e}") from e

    # 批量添加节点ID
    g.add_nodes_from(node_ids)

    # 批量设置节点属性
    nx.set_node_attributes(g, node_kinds, 'kind')
    nx.set_node_attributes(g, node_tokens, 'token')
    nx.set_node_attributes(g, node_diff_tokens, 'diffnochangeToken')
    # --- 优化点1结束 ---

    # --- 优化点2: 批量添加边 ---
    # 准备边数据，格式为 [(source, target), ...]
    # 假设 ast_feature["edges"] 中的每个元素都是 {"S": ..., "T": ...} 字典
    # 同样增加健壮性检查
    edges_data = []
    try:
        for i, e in enumerate(ast_feature["edges"]):
            if not isinstance(e, dict) or "S" not in e or "T" not in e:
                raise KeyError(f"边列表中第 {i} 个元素缺少 'S' 或 'T' 键或不是字典: {e}")
            edges_data.append((e["S"], e["T"]))
    except KeyError as e:
        raise KeyError(f"处理边时出错: {e}") from e

    # 批量添加边
    g.add_edges_from(edges_data)
    # --- 优化点2结束 ---

    # 3. 返回构建好的图
    return g


def cosine_similarity(vector_a, vector_b):
    if len(vector_a) != len(vector_b):
        return -1
    dot_product = np.dot(vector_a, vector_b)  # 向量点乘
    norm_a = np.linalg.norm(vector_a)         # A向量的模
    norm_b = np.linalg.norm(vector_b)         # B向量的模
    if norm_a == 0 or norm_b == 0:
        return 0
    similarity = dot_product / (norm_a * norm_b)  # 余弦相似度公式
    return similarity


def get_mongo_collection(collection_name):
    client = MongoClient(Config.MONGO_URI)
    db = client[Config.DB_NAME]
    return db[collection_name]


def jaccard_similarity(set1, set2):
    if not isinstance(set1, set) or not isinstance(set2, set):
        raise TypeError("Both inputs must be sets")
    intersection = set1 & set2
    union = set1 | set2
    if len(union) == 0:
        return 0.0
    return len(intersection) / len(union)


def calculate_node_similarity(g1, g2, n1, n2):
    """计算两个节点内容的相似度（基于token）"""
    token1 = set(g1.nodes[n1].get("token", []))
    token2 = set(g2.nodes[n2].get("token", []))
    return jaccard_similarity(token1, token2)


def calculate_subtree_similarity(g1, g2, mapping):
    """计算两棵子树的内容相似度（基于映射节点对的token）"""
    if not mapping:
        return 0.0

    similarities = []
    for n1, n2 in mapping.items():
        sim = calculate_node_similarity(g1, g2, n1, n2)
        similarities.append(sim)

    return sum(similarities) / len(similarities) if similarities else 0.0


def convert_objectid(obj):
    """递归转换对象中的所有ObjectId为字符串"""
    if isinstance(obj, ObjectId):
        return str(obj)
    elif isinstance(obj, dict):
        return {k: convert_objectid(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_objectid(item) for item in obj]
    else:
        return obj